import itertools
from collections import defaultdict

def itemset_issubset_transaction(itemset, transaction):
    if isinstance(transaction, (set, frozenset)):
        return transaction.issuperset(itemset)
    else:
        return set(itemset).issubset(transaction)

def __remove_itemsets_not_having_required_support(itemsetToSupportTable, support):
    fisDontHaveRequiredSupport = [item for item, count in itemsetToSupportTable.iteritems() \
            if count < support]
    for fi in fisDontHaveRequiredSupport:
        del itemsetToSupportTable[fi]

def build_size_1_frequent_itemsets(transactions, support=2): # returns a dict, frequent itemset -> support
    itemsetToSupportTable = defaultdict(int)
    for item in itertools.chain(*transactions):
        itemsetToSupportTable[( item, )] += 1
    
    __remove_itemsets_not_having_required_support(itemsetToSupportTable, support)
    
    return itemsetToSupportTable

def __build_next_size_frequent_itemsets_iter(itemsetSupportIter, transactions, support=2,
        hintData=None): # returns a dict, frequent itemset -> support
    itemsets = sorted(map(lambda i_s: i_s[0], itemsetSupportIter))
    candidateItemsets = [] # commonHeader, list of (item1, item2)
    for k, g in itertools.groupby(itemsets, key=lambda i: i[:-1]):
        lg = list(g)
        #assert len(set(lg)) == len(lg)
        tl = []; candidateItemsets.append(( tuple(k), tl ))
        for i, itemset1 in enumerate(lg):
            itemset1Tail = itemset1[-1]
            tl.extend(( itemset1Tail, itemset2[-1] ) for itemset2 in lg[i + 1:])
    
    tis = list(range(len(transactions))) if hintData is None else hintData
    referedTis = []
    candidateItemsetToSupportTable = defaultdict(lambda: defaultdict(int))
    for ti in tis:
        t = transactions[ti]
        for commonHeader, tailList in candidateItemsets:
            if not itemset_issubset_transaction(commonHeader, t): continue
            for tail in tailList:
                if itemset_issubset_transaction(tail, t):
                    candidateItemsetToSupportTable[commonHeader][tail] += 1
                    referedTis.append(ti)
    if hintData is not None:
        hintData[:] = sorted(set(referedTis))
   
    for commonHeader, tailTable in candidateItemsetToSupportTable.iteritems():
        __remove_itemsets_not_having_required_support(tailTable, support)
    
    for commonHeader, tailTable in sorted(candidateItemsetToSupportTable.iteritems()):
        for tail, count in sorted(tailTable.iteritems()):
            yield commonHeader + tail, count

def build_next_size_frequent_itemsets(itemsetToSupportTable, transactions, support=2,
        hintData=None): # returns a dict, frequent itemset -> support
    return dict(__build_next_size_frequent_itemsets_iter(itemsetToSupportTable.iteritems(), 
            transactions, support, hintData))

def build_size_n_frequent_itemsets_depthfirst(size, transactions, support=2, log_write=None): # returns a dict, frequent itemset -> support
    assert size >= 1
    itemsetToSupportTable = build_size_1_frequent_itemsets(transactions, support)
    
    resultTable = dict()
    def b_i(curSize, itemsetSupportIter, curHintData):
        if curSize >= size:
            resultTable.update(itemsetSupportIter)
            return
        for _, g in itertools.groupby( \
                __build_next_size_frequent_itemsets_iter(itemsetSupportIter, transactions, support, curHintData), 
                key=lambda k_v: k_v[0][:-1]):
            partialItemsetAndSupports = list(g)
            if partialItemsetAndSupports: b_i(curSize + 1, partialItemsetAndSupports, list(curHintData))
    hintData = list(xrange(len(transactions)))
    b_i(1, itemsetToSupportTable.iteritems(), hintData)
    return resultTable

def build_size_n_frequent_itemsets(size, transactions, support=2, log_write=None): # returns a dict, frequent itemset -> support
    assert size >= 1
    log_write and log_write("log> building itemsets of size: %d\n" % 1)
    itemsetToSupportTable = build_size_1_frequent_itemsets(transactions, support)
    
    hintData = set(xrange(len(transactions)))
    for curSize in xrange(1, size):
        log_write and log_write("log> building itemsets of size: %d\n" % (curSize + 1))
        if not itemsetToSupportTable: break # for curSize
        itemsetToSupportTable = build_next_size_frequent_itemsets(itemsetToSupportTable, transactions, support,
            hintData)
    return itemsetToSupportTable

def calc_itemset_to_supporting_transactions_table(itemsets, transactions): # returns a dict, itemset -> supporting transactions
    itemsetToItemsetSet = dict(( i, set(i) ) for i in itemsets)
    itemsetToSupportingTransactionInicesTable = dict(( i, [] ) for i in itemsets)
    for ti, t in enumerate(transactions):
        for itemset, itemsetSet in itemsetToItemsetSet.iteritems():
            if itemsetSet.issubset(t):
                itemsetToSupportingTransactionInicesTable[itemset].append(ti)
    return itemsetToSupportingTransactionInicesTable

if __name__ == '__main__':
    transactions = [ 
        ( 'a', 'b', 'c' ),
        ( 'a', 'b' ),
        ( 'a', 'b', 'd' ),
        ( 'b', 'd' )
    ]
    
    fi2sp = build_size_1_frequent_itemsets(transactions, 2)
    while fi2sp:
        #for fi, s in sorted(fi2sp.iteritems()):
        #    print fi, s
        fi2t = calc_itemset_to_supporting_transactions_table(fi2sp.keys(), transactions)
        for itemset, supportingTransactionIndices in sorted(fi2t.iteritems()):
            assert len(supportingTransactionIndices) == fi2sp[itemset]
            print itemset, supportingTransactionIndices
        fi2sp = build_next_size_frequent_itemsets(fi2sp, transactions, 2)
    
    